#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <fcntl.h>
#include <zlib.h>
#include <assert.h>
#include "buffer.h"
#include "flv.h"
#include "screen.h"

#define BLOCK_DIRTY 1
struct FLVAPI RGB {
	uc r;
	uc g;
	uc b;
};
FLVAPI void RGB2GBR(struct RGB *src, struct BGR *dest) {
	dest->r = src->r;
	dest->g = src->g;
	dest->b = src->b;
}

FLVAPI int compareRGB_BGR(struct RGB *p1, struct BGR *p2) 
{
	if(p1->r != p2->r || p1->g != p2->g ||p1->b != p2->b ) return 1;
	return 0;
}
FLVAPI void markBlock(ScreenVideo *video, int x, int y) 
{
	int col = x / video->blockSize;
	int row = y / video->blockSize;
	video->blockMap[row * video->cols + col] = BLOCK_DIRTY;
}	
FLVAPI void resetBlockMap(ScreenVideo *video) 
{
	memset(video->blockMap, 0, video->cols * video->rows);
}
FLVAPI void markAllBlocks(ScreenVideo *video) 
{
	memset(video->blockMap, 1, video->cols * video->rows);
}
void __freescreen(Stream *stream)
{
	ScreenVideo *v;
	if(stream == NULL) return;
	v = (ScreenVideo *)stream;
	if(v->image != NULL) 
		free(v->image);
	if(v->blockMap != NULL)
		free(v->blockMap);
	free(v);
}
FLVAPI int getImageBlock(ScreenVideo *video, int row, int col, struct BGR *block) 
{
	int blockSize = video->blockSize;
	int x_l, x_h, y_l, y_h;
	int x, y, line, pos;
	int size = 0;
	x_l = col * blockSize;
	y_l = row * blockSize;
	if(col == video->cols - 1) 
		x_h = video->imageWidth;
	else
		x_h = x_l + blockSize;
	if(row == video->rows - 1)
		y_h = video->imageHeight;
	else
		y_h = y_l + blockSize;
	for(y = y_l; y < y_h; y++) {
		line = y * video->imageWidth;
		for(x = x_l; x < x_h; x++) 
		{
			pos = line + x;
			block[size++] = video->image[pos];
		}
	}
	return size;
}
FLVAPI int writeEmptyImageBlock(ScreenVideo *video, CBuffer *buf) 
{
	return buf->WriteUI16(0);
}
FLVAPI int writeImageBlock(ScreenVideo* video, CBuffer *buf, int row, int col) 
{
	int cnt=video->blockSize * video->blockSize;
	BGR* src=new BGR[cnt];
	uc* dest=new uc[video->blockSize * video->blockSize * 4];
	int size, ret;
	unsigned long zSize;
	zSize = video->blockSize * video->blockSize * 4;
	size = getImageBlock(video, row, col, src);
	ret = compress2((Bytef *)dest, &zSize, (Bytef *)src, size * sizeof(struct BGR),video->compression);
	if(ret != Z_OK) 
	{
		printf("compression failed %i\n", ret);
		return -1;
	}
	buf->WriteUI16(zSize);
	buf->Write(dest, zSize);
	return zSize + 2;
}
FLVAPI int writeVideoPacketHeader(ScreenVideo *video, CBuffer *buf) 
{
	uc c;
	uc blockSize;
	int ret = 0;
	blockSize = (video->blockSize / 16) - 1;
	/* UB[4] blockWidth */
	c = (blockSize & 0x0f) << 4;
	/* UB[12] imageWidth */
	c |= ( (video->imageWidth & 0x0f00) >> 8 ) & 0x0f;
	ret += buf->Writec(c);
	ret += buf->Writec(video->imageWidth);
	/* UB[4] blockHeight */
	c = (blockSize & 0x0f) << 4;  
	/* UB[12] imageHeight */
	c |= ( (video->imageHeight & 0x0f00) >> 8 ) & 0x0f;
	ret += buf->Writec(c);
	ret += buf->Writec(video->imageHeight);
	return ret;
}
FLVAPI int writeVideoData(ScreenVideo *video, CBuffer *buf) 
{
	int row, col, pos, ret;
	ret = writeVideoPacketHeader(video, buf);	
	for(row = 0; row < video->rows; row++) 
	{
		for(col = 0; col < video->cols; col++) 
		{
			pos = row * video->cols + col;
			if(video->blockMap[pos] == BLOCK_DIRTY) 
			{
				ret += writeImageBlock(video, buf, row, col);
			}
			else
				ret += writeEmptyImageBlock(video, buf);
		}
	}
	return ret;
}
FLVAPI int markBlocks(ScreenVideo *video,PixelData *pixelData) 
{
	int rowstride;
	int x, y, row = 0;
	struct RGB *rgb;
	struct BGR *bgr;
	int size = 0;
	rowstride = pixelData->width * pixelData->n_channels + pixelData->rowPadding;
	if(video->frame == 0)
	{
		markAllBlocks(video);
		return 1;
	}
	resetBlockMap(video);
	for( y = 0; y < pixelData->height; y++) 
	{
		if(pixelData->rowOrder == TOPDOWN)
			row = (pixelData->height - y - 1) * rowstride;
		else if (pixelData->rowOrder == BOTTOMUP) 
			row = y * rowstride;
		for( x = 0; x < pixelData->width; x++) 
		{
			int pos = row + x * pixelData->n_channels;
			rgb = (struct RGB*)(pixelData->data + pos);
			bgr = video->image + size++;
			if(compareRGB_BGR(rgb, bgr)) 
			{
				RGB2GBR(rgb, bgr);
				markBlock(video, x, y);
			}
		}
	}
	return size;
}
FLVAPI int writeTag(ScreenVideo *video, CBuffer *buf)
{	
	Stream *stream = (Stream *)video;	
	if(video->frame == 0)
		stream->hdr.video.frameType = FLV_FRAME_KEY;
	else
		stream->hdr.video.frameType = FLV_FRAME_INTER;
	return writeVideoData(video, buf);
}
int writeDataScreen(Stream *stream, CBuffer *buf, int *advance,void *data) 
{	
	ScreenVideo *video;
	int size;
	struct PixelData *pixelData;	
	if(stream == NULL || data == NULL || buf == NULL) return -E_STREAM;
	video = (ScreenVideo *)stream;
	pixelData = (struct PixelData *)data;
	if(video->imageWidth != pixelData->width || video->imageHeight != pixelData->height) 
	{
		printf("resize not supported\n");
		return -E_STREAM;
	}
	markBlocks(video, pixelData);	
	size = writeTag(video, buf);
	video->frame++;
	return size;	
}
void reset(Stream *stream)
{
	ScreenVideo *video = (ScreenVideo *)stream;
	video->frame = 0;
	resetBlockMap(video);
}

CScreenVideo* CScreenVideo::Create(int width, int height, int blockSize, int comp) 
{
	CScreenVideo *video;
	CStream *stream;
	int imageSize;	
	if(blockSize < 16 || blockSize > 256) 
	{
		printf("blocksize must lie between 16 and 256\n");
		return NULL;
	}
	video =new CScreenVideo;
	if(video == NULL) return NULL;
	stream = &video->m_screenvideo.stream;
	stream->m_stream.writeData = writeDataScreen;
	stream->m_stream.readData = NULL;
	stream->m_stream.freeStream = __freescreen;
	stream->m_stream.resetStream = reset;
	stream->m_stream.type = FLV_TAG_TYPE_VIDEO;
	stream->m_stream.hdr.video.frameType = FLV_FRAME_KEY;
	stream->m_stream.hdr.video.codec = FLV_CODEC_SCREEN;
	video->m_screenvideo.compression = comp;
	video->m_screenvideo.imageWidth = width;
	video->m_screenvideo.imageHeight = height;
	video->m_screenvideo.blockSize = blockSize;
	video->m_screenvideo.rows = height / blockSize;
	if(height % blockSize)
		video->m_screenvideo.rows++;
	video->m_screenvideo.cols = width / blockSize;
	if(width % blockSize)
		video->m_screenvideo.cols++;
	video->m_screenvideo.blockMap = (char *)malloc(video->m_screenvideo.cols * video->m_screenvideo.rows);
	if(video->m_screenvideo.blockMap == NULL)
	{
		free(video);
		return NULL;
	}
	imageSize = sizeof(struct BGR) * width * height;
	video->m_screenvideo.image = (struct BGR *)malloc(imageSize);
	if(video->m_screenvideo.image == NULL)
	{
		free(video->m_screenvideo.blockMap);
		free(video);
		return NULL;
	}
	memset(video->m_screenvideo.image, 0, imageSize);	
	stream->Rewind();
	video->m_screenvideo.stream=*stream;
	return video;
}